#!/usr/bin/python -O

import os, sys
from datetime import datetime
import re


today_str = datetime.today().strftime('%Y%m%d')
rx_stockinfo = re.compile('(.+)="(.+)"')

def get_stockinfo(code):
    import requests
    resp = requests.get('http://hq.sinajs.cn/list=%s' % code)

    lines = resp.text.split('\n')

    m = rx_stockinfo.search(lines[0])

    if m:
        items = m.group(2).split(',')
        info = {
             'code'   : m.group(1)[-6:],
             'open'   : float(items[1]),
             'close'  : float(items[3]),
             'last'   : float(items[2]),
             'inc'    : float(items[3]) - float(items[2]),
             'low'    : float(items[5]),
             'high'   : float(items[4]),
             'volumn' : int(float(items[8]) / 100),
             'amount' : int(float(items[9]) / 10000) }

        if info['close'] == 0 or info['close'] == info['last']:
            info['inc_p'] = 0
        else:
            info['inc_p'] = (info['close'] - info['last']) / info['last'] * 100

        return info


def save_stock(file_handle, stockinfo):
    if __debug__:
        print stockinfo

    file_handle.write(','.join(['S{0}'.format(stockinfo['code']),
                                str(stockinfo['open']),
                                str(stockinfo['close']),
                                str(stockinfo['last']),
                                str(stockinfo['inc']),
                                '{0:.3f}'.format(stockinfo['inc_p']),
                                str(stockinfo['low']),
                                str(stockinfo['high']),
                                str(stockinfo['volumn']),
                                str(stockinfo['amount'])]) + '\n')




# -------- main --------
if __name__ == '__main__':

    from optparse import OptionParser

    parser = OptionParser()
    parser.add_option('--mongo',
           help='address of mongodb server')
    parser.add_option('--datadir',
           default='{0}/Copy/stockdata'.format(os.getenv('HOME')),
           help='directory to store stockdata')

    (opts, args) = parser.parse_args()

    codelist = [ 'sz%06d' % i for i in range(1, 1000) + range(2001, 3000)
                                     + range(300001, 301000) ] \
             + [ 'sh%06d' % i for i in range(600000, 602000) ]

    fname = '{0}/{1}.csv'.format(opts.datadir, today_str)

    with file(fname, 'w') as stock_file:
        stock_file.write(','.join(['code',
                                'open',
                                'close',
                                'last',
                                'inc',
                                'inc_p',
                                'low',
                                'high',
                                'volumn',
                                'amount']) + '\n')

        try:
            for code in codelist:
                stockinfo = get_stockinfo(code)
                if stockinfo:
                    save_stock(stock_file, stockinfo)
        except KeyboardInterrupt:
            pass

